#!/usr/bin/env python
import sys
import os.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import simpy
import numpy as np
import pandas as pd
from scipy.stats import norm, entropy, skew, kurtosis
from jmarkov.phase.fit.moments_ctph2 import moments_ctph2
from jmarkov.phase.ctph import ctph
from jmarkov.queue.phph1 import phph1
import itertools
from scipy.interpolate import interp1d
import time
import math





def estimate_rho_from_corr_matrix(corr_matrix):
    n = corr_matrix.shape[0]
    upper_tri = corr_matrix[np.triu_indices(n, k=1)]
    return np.mean(upper_tri)

def generate_batch_correlated_exponentials(num_jobs, n, mean, rho):
    np.random.seed(123)
    cov = rho * np.ones((n, n)) + (1 - rho) * np.eye(n)

    mv_norm = np.random.multivariate_normal(
        mean=np.zeros(n),
        cov=cov,
        size=num_jobs
    )

    uniforms = norm.cdf(mv_norm)

    exponentials = -mean * np.log(uniforms)

    return exponentials
def simulate_queue(N,
                   INTER_ARRIVAL_TIME,
                   TASK_MEAN_DURATION,
                   RHO,
                   SIM_TIME):

    env = simpy.Environment()
    server = simpy.Resource(env, capacity=1)

    waiting_times = []
    system_times = []

    #np.random.seed(123)

    # --------------------------------------------------

    # --------------------------------------------------

    expected_jobs = int(
    np.ceil(2 * SIM_TIME / INTER_ARRIVAL_TIME)
)

    task_pool = generate_batch_correlated_exponentials(
    expected_jobs,
    N,
    TASK_MEAN_DURATION,
    RHO
)


    arrival_gaps = np.random.exponential(
    INTER_ARRIVAL_TIME,
    size=expected_jobs
)

    # --------------------------------------------------

    # --------------------------------------------------

    def job(env, job_id, server):
        arrival_time = env.now
        durations = task_pool[job_id]
        with server.request() as req:
            yield req
            start_service = env.now
            waiting_times.append(
                start_service - arrival_time
            )
            service_time = durations.max()
            yield env.timeout(service_time)
            completion_time = env.now
            system_times.append(
                completion_time - arrival_time
            )

    def job_generator(env, server):
        job_id = 0
        for gap in arrival_gaps:
            yield env.timeout(gap)
            if env.now > SIM_TIME:
                break
            env.process(
                job(env, job_id, server)
            )
            job_id += 1

    env.process(
        job_generator(env, server)
    )

    env.run(until=SIM_TIME)
    num_jobs_used = len(system_times)
    task_matrix = task_pool[:num_jobs_used]

    return (
        task_matrix,
        np.array(waiting_times),
        np.array(system_times)
    )


def fit_phase_type_and_compare(task_matrix, waiting_times, system_times, N, INTER_ARRIVAL_TIME,delta):
    start_time = time.time()
    empirical_corr_matrix = np.corrcoef(task_matrix.T)
    estimated_rho = estimate_rho_from_corr_matrix(empirical_corr_matrix)

    mean_all_tasks = task_matrix.mean()
    l1 = mean_all_tasks
    l2 = task_matrix.var() + l1*l1
    l12 = estimated_rho


    #p = 1-np.sqrt(l12)
    p = 1-l12
    #b = np.sqrt(l2 - l1*l1)
    b = np.sqrt((l2 - l1*l1)/(2-p))
    m1 = l1 - b*p
    gamma = 1 / m1
    #m2 = l2 - l1*l1 + m1*m1 - b*b*(2*p-p*p)
    #fitter = moments_ctph2(m1, m2)


    #PH = fitter.get_ph()

    lda = 1 / INTER_ARRIVAL_TIME
    alpha = np.array([1])
    T = np.array([[-lda]])
    IAT = ctph(alpha, T)

    mu1 = 1/b
    num_tasks = N
    #beta = np.zeros(num_tasks + 2)
    beta = np.zeros(num_tasks + 1)
    #beta[0], beta[1] = PH.alpha[0], PH.alpha[1]
    beta[0] = 1
    #S = np.zeros((num_tasks + 2, num_tasks + 2))
    S = np.zeros((num_tasks + 1, num_tasks + 1))
    #S[0:2, 0:2] = PH.T
    S[0, 0] = -gamma 
    #S[0:2, 2:3] = -PH.T @ np.ones((2, 1))
    S[0, 1] = gamma*p
    for i in range(num_tasks - 1):
        S[i+1, i+1] = -(num_tasks - i) * mu1
        S[i+1, i+2] = (num_tasks - i) * mu1
    #S[num_tasks+1, num_tasks+1] = -mu1
    S[num_tasks, num_tasks] = -mu1
    ST = ctph(beta, S)
    q = phph1(IAT, ST)
    print("beta:")
    print(beta)
    print("S:")
    print(S)
    #print("PH.T:")
    #print(PH.T)
    print("mu1:")
    print(mu1)



    WT = q.wait_time_dist()

    RT = q.resp_time_dist()
    meanRT = RT.mean()
    m2RT = RT.moment(2)
    stdRT = np.sqrt(m2RT - meanRT**2)
    m3RT = RT.moment(3)
    m3RTc = m3RT - 3*meanRT*stdRT**2 - meanRT**3
    skewRT = m3RTc / (stdRT ** 3) 
    m4RT = RT.moment(4) 
    m4RTc = m4RT - 4*meanRT*m3RT + 6*meanRT**2*m2RT - 3*meanRT**4
    kurtRT = m4RTc / (stdRT ** 4) -3
    print("RT.mean():", meanRT)
    print("RT.std():", stdRT)
    print("RT.skew():", skewRT)
    print("RT.kurtosis():", kurtRT)
    cur_time = time.time()
    print(f"Analytical Response time: {cur_time - start_time}")


    # Compare empirical vs theoretical CDFs
    t_max = system_times.max()
    #dt=0.05
    delta2=np.ceil(len(system_times)*delta)
    dt = (t_max-system_times.min())/delta2
    ts = np.arange(system_times.min(), t_max, dt)
    print(f"len(ts): {len(ts)}")

    start_time = time.time()
    #theo_resp_pdf = np.array([RT.pdf([t+dt/2],unif=True) for t in ts])
    theo_resp_pdf = np.array(RT.pdf(ts, unif=True))
    print(f"theo_resp_pdf[0]: {theo_resp_pdf[0]}")
    #print(f"len(theo_resp_pdf): {theo_resp_pdf.shape}")
    theo_resp_pdf[theo_resp_pdf < 0] = 0
    cur_time = time.time()
    print(f"Analytical Response time PDF: {cur_time - start_time}")


    sim_response_times = system_times
    counts_r, edges_r = np.histogram(sim_response_times, bins=len(ts), density=True)
    #print("Teórica")
    #print((theo_resp_pdf))
    #print("Simulada")
    #print((counts_r))

    kl_div=entropy(counts_r + 1e-100, theo_resp_pdf + 1e-100)

    pdf = np.array(theo_resp_pdf)
    pdf = pdf / (pdf.sum() * dt)

    theo_r2 = np.cumsum(pdf) * dt

    cdf = np.clip(theo_r2, 0, 1)
    cdf /= cdf[-1]

    inv_cdf = interp1d(
        cdf,
        ts,
        kind='linear',
        fill_value="extrapolate",
        assume_sorted=True
    )

    # 3. Simulate N samples
    u = np.random.uniform(size=10000)
    samples = inv_cdf(u)

    result = {
            'KL': kl_div,
            'MeanSim': np.mean(system_times),
            'StdSim': np.std(system_times),
            'SkewSim': skew(system_times),
            'KurtosisSim': kurtosis(system_times),
            'MeanTheo': np.mean(samples),
            'StdTheo': np.std(samples),
            'SkewTheo': skew(samples),
            'KurtosisTheo': kurtosis(samples),
            "sim_r":sim_response_times,
            "theo_r":theo_resp_pdf
            #'hist': cdf_values,
            #'theo_pdf': theo_wait_cdf
    }

    return result





#N_values=[5]
#utilizacion=[0.1]
#TASK_MEAN_VALUES=[1]
#RHO_VALUES=[0.1]
#SIM_TIME_VALUES = [10000]


#N_values = [5,10,50,100,500,1000]
#utilizacion=[0.1,0.3,0.5,0.7,0.9]
#RHO_VALUES = [0.3,0.5,0.7,0.9]
#TASK_MEAN_VALUES = [1]
#SIM_TIME_VALUES=[1000]

#N_values = [5,10,50,100,500]
#utilizacion=[0.1,0.3,0.5,0.7,0.9]
#RHO_VALUES = [0.3,0.5,0.7,0.9]
#TASK_MEAN_VALUES = [1]
#SIM_TIME_VALUES=[100000]


N_values=[500]
utilizacion=[0.9]
TASK_MEAN_VALUES=[1]
RHO_VALUES=[0.9]
SIM_TIME_VALUES = [1000000]



delta=0.05
results = []

for N, ut, TM, RHO, SIM_T in itertools.product(
        N_values, utilizacion, TASK_MEAN_VALUES, RHO_VALUES, SIM_TIME_VALUES):
    HN=sum(1/i for i in range(1, N+1))
    IAT=1/(ut/HN)

    print(f"Running N={N}, utilizacion={ut}, RHO={RHO}")

    start_time2 = time.time()


    start_time = time.time()
    task_matrix, waits, systems = simulate_queue(N, IAT, TM, RHO, SIM_T)
    tiempo_sim=time.time() - start_time
    print("Sim time: %s seconds" % (tiempo_sim))


    start_time = time.time()
    #wait_m, resp_m= fit_phase_type_and_compare(task_matrix, waits, systems, N, IAT)
    resp_m= fit_phase_type_and_compare(task_matrix, waits, systems, N, IAT,delta)
    tiempo_est=time.time() - start_time
    print("Est time: %s seconds" % (tiempo_est))

    results.append({
        'N': N,
        'InterArrival': IAT,
        'TaskMean': TM,
        'Rho': RHO,
        'SimTime': SIM_T,
        #'KL_wait': wait_m['KL'],
        #'Mean_wait1': wait_m['Mean1'],
        #'Mean_wait2': wait_m['Mean2'],
        #'Std_wait1': wait_m['Std1'],
        #'Std_wait2': wait_m['Std2'],
        #'Skew_wait1': wait_m['Skew1'],
        #'Skew_wait2': wait_m['Skew2'],
        #'Kurt_wait1': wait_m['Kurtosis1'],
        #'Kurt_wait2': wait_m['Kurtosis2'],
        'KL_resp': resp_m['KL'],
        'Mean_respSim': resp_m['MeanSim'],
        'Mean_respTheo': resp_m['MeanTheo'],
        'Std_respSim': resp_m['StdSim'],
        'Std_respTheo': resp_m['StdTheo'],
        'Skew_respSim': resp_m['SkewSim'],
        'Skew_respTheo': resp_m['SkewTheo'],
        'Kurt_respSim': resp_m['KurtosisSim'],
        'Kurt_respTheo': resp_m['KurtosisTheo'],
        #"counts_w": wait_m['hist'],
        #"theo_w": wait_m['theo_pdf'],
        #"counts_r": resp_m['hist'],
        #"theo_r": resp_m['theo_pdf'],
        "tiempo_sim": tiempo_sim,
        "tiempo_est": tiempo_est,
        "sim_r":resp_m["sim_r"],
        "theo_r":resp_m["theo_r"]
    })

df = pd.DataFrame(results)
print(df)

